#include <iostream>
#include <stdio.h>

extern "C" {
#include "libavutil/log.h"
#include "libavformat/avformat.h"
}

#define ERROR_STRING_SIZE (1024)
#define ADTS_HEADER_LEN (7)

const int samplingFrequencies[] = {
    96000, //0x00
    88200,
    64000,
    48000,
    44100,
    32000, //0x05
    24000,
    22050,
    16000,
    12000,
    11025, //0x0a
    8000
};

int makeAdtsHeader(char * const pAdtsHeader, const int dataLength,
               const int profile, const int sampleRate, const int channels) {
    int samplingFrequenciesIndex = 3; //默认使用 48000Hz
    int adtsLen = dataLength + 7;

    int frequenciesSize = sizeof(samplingFrequencies) / sizeof(samplingFrequencies[0]);
    int i = 0;

    for (int i = 0; i < frequenciesSize; i++) {
        if (samplingFrequencies[i] == sampleRate) {
            samplingFrequenciesIndex = i;
            break;
        }
    }

    if (i >= frequenciesSize) {
        std::cerr << "Unsupport samplerate: " << sampleRate << std::endl;
        return -1;
    }

    pAdtsHeader[0] = 0xff;
    pAdtsHeader[1] = 0xf0;
    pAdtsHeader[1] |= (0 << 3);
    pAdtsHeader[1] |= (0 << 1);
    pAdtsHeader[1] |= 1;

    pAdtsHeader[2] = (profile) << 6;
    pAdtsHeader[2] |= (samplingFrequenciesIndex & 0x0f) << 2;
    pAdtsHeader[2] |= (0 << 1);
    pAdtsHeader[2] |= (channels & 0x04) >> 2;

    pAdtsHeader[3] = (channels & 0x03) << 6;
    pAdtsHeader[3] |= (0 << 5);
    pAdtsHeader[3] |= (0 << 4);
    pAdtsHeader[3] |= (0 << 3);
    pAdtsHeader[3] |= (0 << 2);
    pAdtsHeader[3] |= ((adtsLen & 0x1800) >> 11);

    pAdtsHeader[4] = (uint8_t)((adtsLen & 0x7f8) >> 3);
    pAdtsHeader[5] = (uint8_t)((adtsLen & 0x7) << 5);
    pAdtsHeader[5] |= 0x1f;
    pAdtsHeader[6] = 0xfc;

    return 0;
}

int main(int argc, char *argv[])
{
    if (argc != 4) {
        std::cout << "Usage: app input.mp4 out.h264 out.aac" << std::endl;
        return -1;
    }

    std::string inputFileName = std::string(argv[1]);
    std::string h264FileName  = std::string(argv[2]);
    std::string aacFileName   = std::string(argv[3]);
    FILE *h264Fd = NULL, *aacFd = NULL;

    h264Fd = fopen(h264FileName.c_str(), "wb");
    if (h264Fd == NULL) {
        std::cerr << "Failed to open " << h264FileName << std::endl;
        return -1;
    }

    aacFd = fopen(aacFileName.c_str(), "wb");
    if (aacFd == NULL) {
        std::cerr << "Failed to open " << aacFileName << std::endl;
        return -1;
    }

    AVFormatContext *ifmtCtx = NULL;
    int videoIndex = -1;
    int audioIndex = -1;
    AVPacket *packet = NULL;
    int ret = 0;
    char errors[ERROR_STRING_SIZE+1]; //用于缓存解析 ffmpeg api 返回值的错误 string

    ifmtCtx = avformat_alloc_context();
    if (ifmtCtx == NULL) {
        std::cerr << "avformat_alloc_context failed." << std::endl;
        return -1;
    }

    ret = avformat_open_input(&ifmtCtx, inputFileName.c_str(), NULL, NULL);
    if (ret < 0) {
        av_strerror(ret, errors, ERROR_STRING_SIZE);
        std::cerr << "avformat_open_input failed: " << ret << " " << errors << std::endl;
        avformat_close_input(&ifmtCtx);
        return -1;
    }

    videoIndex = av_find_best_stream(ifmtCtx, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0);
    if (videoIndex == -1) {
        std::cerr << "av_find_best_stream videoIndex failed.";
        avformat_close_input(&ifmtCtx);
        return -1;
    }

    audioIndex = av_find_best_stream(ifmtCtx, AVMEDIA_TYPE_AUDIO, -1, -1, NULL, 0);
    if (audioIndex == -1) {
        std::cerr << "av_find_best_stream audioIndex failed." << std::endl;
        avformat_close_input(&ifmtCtx);
        return -1;
    }

    //h264_mp4toannexb
    const AVBitStreamFilter *bsfilter = av_bsf_get_by_name("h264_mp4toannexb");
    if (bsfilter == NULL) {
        avformat_close_input(&ifmtCtx);
        std::cerr << "av_bsf_get_by_name h264_mp4toannexb failed." << std::endl;
        return -1;
    }

    AVBSFContext *bsfCtx = NULL;
    ret = av_bsf_alloc(bsfilter, &bsfCtx);
    if (ret < 0) {
        av_strerror(ret, errors, ERROR_STRING_SIZE);
        std::cerr << "av_bsf_alloc failed: " << errors << std::endl;
        avformat_close_input(&ifmtCtx);
        return -1;
    }

    ret = avcodec_parameters_copy(bsfCtx->par_in, ifmtCtx->streams[videoIndex]->codecpar);
    if (ret < 0) {
        av_strerror(ret, errors, ERROR_STRING_SIZE);
        std::cerr << "avcodec_parameters_copy failed: " << errors << std::endl;
        avformat_close_input(&ifmtCtx);
        av_bsf_free(&bsfCtx);
        return -1;
    }

    ret = av_bsf_init(bsfCtx);
    if (ret < 0) {
        av_strerror(ret, errors, ERROR_STRING_SIZE);
        std::cerr << "av_bsf_init failed: " << errors << std::endl;
        avformat_close_input(&ifmtCtx);
        av_bsf_free(&bsfCtx);
        return -1;
    }

    packet = av_packet_alloc();
    av_init_packet(packet);

    while (1) {
        ret = av_read_frame(ifmtCtx, packet); //不会释放packet的buf，如果外部不去释放，会内存泄露
        if (ret < 0) {
            av_strerror(ret, errors, ERROR_STRING_SIZE);
            std::cerr << "av_read_frame failed: " << errors << std::endl;
            break;
        }

        //av_read_frame 成功读取到 packet，则外部需要进行 buf 释放
        if (packet->stream_index == videoIndex) {
            ret = av_bsf_send_packet(bsfCtx, packet); //内部把buf转移到自己bsf内部
            if (ret < 0) { //基本不会进入该逻辑
                av_strerror(ret, errors, ERROR_STRING_SIZE);
                std::cerr << "av_bsf_send_packet failed: " << errors << std::endl;
                av_packet_unref(packet);
                continue;
            }

            while (1) {
                ret = av_bsf_receive_packet(bsfCtx, packet);
                if (ret != 0)
                    break;

                size_t size = fwrite(packet->data, 1, packet->size, h264Fd);
                if (size != packet->size)
                    av_log(NULL, AV_LOG_DEBUG, "h264 warning, length of writed data isn't equal packet->size(%d, %d)\n", size, packet->size);

                av_packet_unref(packet);
            }
        }
        else if (packet->stream_index == audioIndex) {
            char adtsHeaderBuf[7] = {0};
            makeAdtsHeader(adtsHeaderBuf, packet->size,
                          ifmtCtx->streams[audioIndex]->codecpar->profile,
                          ifmtCtx->streams[audioIndex]->codecpar->sample_rate,
                          ifmtCtx->streams[audioIndex]->codecpar->channels);
            fwrite(adtsHeaderBuf, 1, 7, aacFd); //写 adts header，ts 流不适用，ts 流分离出来的 packet 带了 adts header
            size_t size = fwrite(packet->data, 1, packet->size, aacFd);
            if (size != packet->size)
                av_log(NULL, AV_LOG_DEBUG, "aac warning, legnth of writed data isn't equal packet.size(%d, %d)\n", size, packet->size);

            av_packet_unref(packet);
        }
        else {
            av_packet_unref(packet);
        }
    }

    std::cout << "While finished" << std::endl;

    if (h264Fd != NULL)
        fclose(h264Fd);

    if (aacFd != NULL)
        fclose(aacFd);

    if (packet != NULL)
        av_packet_free(&packet);

    if (ifmtCtx != NULL)
        avformat_close_input(&ifmtCtx);

    return 0;
}
